//
// Description: 1337. 矩阵中战斗力最弱的 K 行
// Created by Loading on 2021/8/1.
//

#include <bits/stdc++.h>

using namespace std;

vector<int> kWeakestRows(vector<vector<int>> &mat, int k) {
    vector<pair<int, int>> effect;
    for (int i = 0; i < mat.size(); ++i) {
        // 二分判定军人数量
        int l = 0, r = (int) mat[i].size() - 1;
        // 最后一位军人位置
        int pos = -1;
        while (l <= r) {
            int mid = l + (r - l) / 2;
            if (mat[i][mid] == 0) {
                r = mid - 1;
            } else {
                pos = mid;
                l = mid + 1;
            }
        }
        effect.emplace_back(pos + 1, i);
    }

    // 优先级队列
    priority_queue q(greater<>(), effect);
    vector<int> res;
    for (int i = 0; i < k; ++i) {
        res.emplace_back(q.top().second);
        q.pop();
    }

    return res;
}

int main() {
    vector<vector<int>> mat = {{1, 1, 0, 0, 0},
                               {1, 1, 1, 1, 0},
                               {1, 0, 0, 0, 0},
                               {1, 1, 0, 0, 0},
                               {1, 1, 1, 1, 1}};
    int k = 3;

    vector<int> res = kWeakestRows(mat, k);
    for (auto &x : res) {
        cout << x << " ";
    }
    cout << endl;
}